import argparse
import json
import os
import fitz

from dots_ocr import DotsOCRParser
from dots_ocr.utils import dict_promptmode_to_prompt

parser = argparse.ArgumentParser()
parser.add_argument("--src", type=str, default="教师用书")
parser.add_argument("--dst", type=str, default="target")
parser.add_argument("--info", type=str, default="ocr_book.json", help="processed pdf files")
parser.add_argument("--force", type=str, default="n", help="y or n, whether to force ocr")
parser.add_argument("--dpi", type=int, default=200)
parser.add_argument("--num_thread", type=int, default=32)
parser.add_argument("--max_completion_tokens", type=int, default=16384)
parser.add_argument("--temperature", type=float, default=0.1)
parser.add_argument("--top_p", type=float, default=1.0)
args = parser.parse_args()


def main():
    file_type = "pdf"
    info_json = {}
    if os.path.exists(args.info):
        with open(args.info, 'r', encoding='utf8') as fin:
            info_json = json.load(fin)
    failed = []
    prompt_mode = "prompt_layout_all_en"
    prompt = dict_promptmode_to_prompt[prompt_mode]
    dots_ocr_parser = DotsOCRParser(
        ip='localhost',
        port=8000,
        model_name="model",
        temperature=args.temperature,
        top_p=args.top_p,
        max_completion_tokens=args.max_completion_tokens,
        num_thread=args.num_thread,
        dpi=args.dpi,
        output_dir=args.dst,
        min_pixels=None,
        max_pixels=None,
        use_hf=True,
    )
    wait_list = []
    for dir_name, _, filenames in os.walk(args.src):
        pre = "-".join(dir_name.split("/")) if "/" in dir_name else "-".join(dir_name.split("\\"))
        for filename in filenames:
            filepath = os.path.join(dir_name, filename)
            if not filepath.endswith(".pdf"):
                continue
            filename = filename.replace(".pdf", "")
            filename = pre + "-" + filename
            filename = filename.replace(" ", "_")
            if filename in info_json:
                print(filename, "-->", "skipped")
                continue
            wait_list.append((filepath, filename))
    num = 0
    for pdf_path, pdf_name in wait_list:
        num += 1
        print("{}/{}".format(num, len(wait_list)))
        dots_ocr_parser.parse_file(
            pdf_path,
            output_dir=args.dst,
            prompt_mode=prompt_mode,
            bbox=None,
            fitz_preprocess=True,
        )
        info_json[pdf_name] = pdf_path
        with open(args.info, 'w', encoding='utf8') as fout:
            json.dump(info_json, fout, indent=2)


if __name__ == '__main__':
    main()
